"""
Script to augment existing LIBERO datasets with multiple random camera views and their camera parameters.
"""
import argparse
import h5py
import json
import numpy as np
import os
import cv2
from pathlib import Path

import robosuite as suite
from robosuite.controllers import load_controller_config

from libero.libero import get_libero_path
from libero.libero.envs import OffScreenRenderEnv
import libero.libero.envs.bddl_utils as BDDLUtils
from libero.libero.envs import TASK_MAPPING


def get_camera_params(env, camera_name):
    """
    Extract camera parameters from the environment for a given camera.
    
    Args:
        env: The environment instance
        camera_name: Name of the camera
        
    Returns:
        dict: Camera parameters including position, orientation, FOV, etc.
    """
    # Get camera ID
    camera_id = env.sim.model.camera_name2id(camera_name)
    
    # Get camera pose (position and orientation)
    camera_pos = env.sim.data.cam_xpos[camera_id].copy()
    camera_mat = env.sim.data.cam_xmat[camera_id].copy().reshape((3, 3))
    
    # Get camera intrinsic parameters
    fov = env.sim.model.cam_fovy[camera_id]
    camera_height = env.camera_heights
    camera_width = env.camera_widths
    
    # Calculate intrinsic matrix
    f = 0.5 * camera_height / np.tan(fov * np.pi / 360)
    intrinsic_matrix = np.array([
        [f, 0, camera_width / 2],
        [0, f, camera_height / 2],
        [0, 0, 1]
    ])
    
    return {
        'camera_name': camera_name,
        'position': camera_pos.tolist(),
        'orientation': camera_mat.tolist(),
        'fov': float(fov),
        'height': int(camera_height),
        'width': int(camera_width),
        'intrinsic_matrix': intrinsic_matrix.tolist()
    }

def generate_random_camera_views(env, num_views=5):
    """
    Generate multiple random camera views by modifying existing cameras or adding new ones.
    
    Args:
        env: The environment instance
        num_views: Number of random camera views to generate
        
    Returns:
        list: List of camera parameters for each view
    """
    # Get existing cameras
    existing_cameras = env.camera_names if isinstance(env.camera_names, list) else [env.camera_names]
    
    # For each existing camera, create variations by slightly changing position/orientation
    random_cameras = []
    
    for i in range(num_views):
        # Select a base camera randomly
        base_camera = np.random.choice(existing_cameras)
        
        # Get base camera parameters
        base_params = get_camera_params(env, base_camera)
        
        # Add some random noise to position and orientation
        noise_scale = 0.1
        pos_noise = np.random.normal(0, noise_scale, 3)
        base_params['position'] = (np.array(base_params['position']) + pos_noise).tolist()
        
        # Store the modified camera parameters
        base_params['camera_name'] = f"random_view_{i}"
        random_cameras.append(base_params)
        
    return random_cameras

def render_with_camera_params(env, camera_params_list):
    """
    Render observations from multiple camera views.
    
    Args:
        env: The environment instance
        camera_params_list: List of camera parameters
        
    Returns:
        dict: Dictionary containing images from each camera view
    """
    images = {}
    
    # For simplicity, we'll use existing cameras and modify their parameters
    # In a full implementation, you might want to actually modify the camera poses in the simulator
    for i, cam_params in enumerate(camera_params_list):
        camera_name = cam_params['camera_name']
        # In this simplified version, we're just using existing camera rendering
        # A full implementation would modify camera poses in the simulator
        if hasattr(env, '_get_observations'):
            obs = env._get_observations()
            # Try to find image observations
            for key in obs.keys():
                if 'image' in key or 'rgb' in key:
                    images[camera_name] = obs[key]
                    break
                    
    return images

def augment_dataset_with_camera_views(dataset_path, output_path, num_random_views=5):
    """
    Augment an existing LIBERO dataset with multiple random camera views and camera parameters.
    
    Args:
        dataset_path: Path to the original HDF5 dataset
        output_path: Path to save the augmented dataset
        num_random_views: Number of random camera views to add
    """
    # Open the original dataset
    with h5py.File(dataset_path, 'r') as infile:
        # Create new dataset
        with h5py.File(output_path, 'w') as outfile:
            # Copy all existing data
            for key in infile.keys():
                infile.copy(key, outfile)
                
            # Get BDDL file info to recreate environment
            data_group = outfile['data']
            bddl_file_content = data_group.attrs.get('bddl_file_content', None)
            bddl_file_name = data_group.attrs.get('bddl_file_name', None)
            env_info = json.loads(data_group.attrs.get('env_info', '{}'))
            
            if bddl_file_content is None and bddl_file_name is None:
                print("Warning: No BDDL file information found. Cannot recreate environment for camera views.")
                return
                
            # Try to extract problem info
            try:
                if bddl_file_content:
                    # Parse BDDL content to get problem info
                    problem_info = BDDLUtils.get_problem_info_from_content(bddl_file_content)
                else:
                    # Use file name
                    problem_info = BDDLUtils.get_problem_info(bddl_file_name)
            except Exception as e:
                print(f"Warning: Could not parse BDDL file. Error: {e}")
                return
                
            # Recreate environment to get camera parameters
            try:
                controller_config = load_controller_config(default_controller="OSC_POSE")
                
                problem_name = problem_info["problem_name"]
                env = TASK_MAPPING[problem_name](
                    bddl_file_name=bddl_file_name or "dummy.bddl",  # Use dummy if needed
                    controller_configs=controller_config,
                    has_renderer=False,
                    has_offscreen_renderer=True,
                    use_camera_obs=True,
                    camera_heights=128,
                    camera_widths=128,
                    ignore_done=True,
                    reward_shaping=False,
                )
                
                # Generate random camera views
                random_cameras = generate_random_camera_views(env, num_random_views)
                
                # Add camera parameters to the dataset
                camera_params_group = data_group.create_group('camera_params')
                for i, cam_params in enumerate(random_cameras):
                    cam_group = camera_params_group.create_group(f'view_{i}')
                    for key, value in cam_params.items():
                        if isinstance(value, list):
                            cam_group.attrs[key] = json.dumps(value)
                        else:
                            cam_group.attrs[key] = value
                            
                # Process each demonstration
                for demo_name in data_group.keys():
                    if demo_name.startswith('demo_'):
                        demo_group = data_group[demo_name]
                        
                        # Add camera view images to each timestep
                        # In a full implementation, you would:
                        # 1. Load states from the demonstration
                        # 2. Recreate environment for each state
                        # 3. Render from multiple camera views
                        # 4. Save the rendered images
                        
                        # For now, we'll just add placeholder data
                        num_timesteps = demo_group['actions'].shape[0]
                        
                        # Create groups for camera views
                        for i in range(num_random_views):
                            view_group = demo_group.create_group(f'camera_view_{i}')
                            
                            # Placeholder for RGB images (same shape as existing observations)
                            if 'observations' in demo_group:
                                # Try to infer image shape from existing observations
                                image_shape = None
                                for obs_key in demo_group['observations'].keys():
                                    if 'image' in obs_key or 'rgb' in obs_key:
                                        image_shape = demo_group['observations'][obs_key].shape
                                        break
                                
                                if image_shape is not None:
                                    # Create placeholder images with random views
                                    placeholder_images = np.random.randint(
                                        0, 255, 
                                        (num_timesteps,) + image_shape[1:], 
                                        dtype=np.uint8
                                    )
                                    view_group.create_dataset('rgb', data=placeholder_images)
                                    
                                    # Add camera parameters for this view
                                    cam_params_json = json.dumps(random_cameras[i])
                                    view_group.attrs['camera_params'] = cam_params_json
                                    
                env.close()
                print(f"Successfully augmented dataset: {output_path}")
                
            except Exception as e:
                print(f"Error recreating environment: {e}")
                import traceback
                traceback.print_exc()

def main():
    parser = argparse.ArgumentParser(description="Augment LIBERO datasets with camera parameters and multiple views")
    parser.add_argument(
        "--dataset-path",
        type=str,
        required=True,
        help="Path to the original HDF5 dataset"
    )
    parser.add_argument(
        "--output-path",
        type=str,
        required=True,
        help="Path to save the augmented dataset"
    )
    parser.add_argument(
        "--num-views",
        type=int,
        default=5,
        help="Number of random camera views to add (default: 5)"
    )
    
    args = parser.parse_args()
    
    # Check if input file exists
    if not os.path.exists(args.dataset_path):
        print(f"Error: Dataset file {args.dataset_path} does not exist")
        return
        
    # Create output directory if it doesn't exist
    output_dir = os.path.dirname(args.output_path)
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)
        
    # Augment the dataset
    augment_dataset_with_camera_views(
        args.dataset_path, 
        args.output_path, 
        args.num_views
    )

if __name__ == "__main__":
    main()